from pymilvus import MilvusClient, DataType, model
from import_excel import get_all_data_from_xls
from openai import OpenAI
from qianwen_embedding import qianwen_embedding


client = MilvusClient(
    uri="http://127.0.0.1:19530",
)

# 使用数据库：my_database_1
client.use_database(db_name="my_database_1")

# 定义collection字段(schema)
schema = client.create_schema() # auto_id=True 表示使用自增id
schema.add_field(field_name="article_id", datatype=DataType.INT64, is_primary=True, description="article_id") # 必须要有主键不然创建会报错
schema.add_field(field_name="user_name", datatype=DataType.VARCHAR, max_length=255, description="专家名称")
schema.add_field(field_name="title", datatype=DataType.VARCHAR, max_length=255, description="文章标题")
schema.add_field(field_name="analyse", datatype=DataType.VARCHAR, max_length=2000, description="文章内容")
schema.add_field(field_name="title_vector", datatype=DataType.FLOAT_VECTOR, dim=512, description="文章标题向量")
schema.add_field(field_name="analyse_vector", datatype=DataType.FLOAT_VECTOR, dim=512, description="文章内容向量")

# 创建索引字段
index_params = client.prepare_index_params()
index_params.add_index(field_name="article_id", index_type="AUTOINDEX")
index_params.add_index(field_name="user_name", index_type="AUTOINDEX")
index_params.add_index(field_name="title", index_type="AUTOINDEX")
index_params.add_index(field_name="title_vector", index_type="HNSW", metric_type="COSINE")
index_params.add_index(field_name="analyse_vector", index_type="HNSW", metric_type="COSINE")

# 先删除集合
client.drop_collection('demo_analyse')

if not client.has_collection('demo_analyse'):
    client.create_collection(
        collection_name="demo_analyse",
        dimension=768, # 向量的维度
        metric_type="COSINE", # 这段代码设置了相似度计算的度量类型为"COSINE"（余弦相似度）。余弦相似度通过计算两个向量夹角的余弦值来衡量它们的相似程度，值域为[-1,1]，常用于文本相似度计算和推荐系统中。
        consistency_level="Bounded", # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/consistency.md#Consistency-Level for more details.
        auto_id=True,
        schema=schema, # 使用定义的schema
        index_params=index_params, # 使用定义的索引字段
    )

# 准备数据
file_path = "分析.xls"  # 替换为你的XLS文件路径
data = get_all_data_from_xls(file_path)

# 文本转成向量
# 安装：pip install "pymilvus[model]"
# If connection to https://huggingface.co/ failed, uncomment the following path
# import os
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# This will download a small embedding model "paraphrase-albert-small-v2" (~50MB).
# embedding_fn = model.DefaultEmbeddingFunction() # 需要魔法才能使用

# 使用千问的embbeding模型

# 格式化要插入的数据
insertData = []
for item in data:
    # print(item)
    # print(qianwen_embedding(item['title']))
    insertData.append({
        "article_id": item['article_id'],
        "user_name": item['user_name'],
        "title": item['title'],
        "analyse": item['analyse'],
        "title_vector": qianwen_embedding(item['title']),
        "analyse_vector": qianwen_embedding(item['analyse']),
    })


# 插入数据
res = client.insert(collection_name="demo_analyse", data=insertData)
print(res)
